
import logging
import os
import sys
import torch
import pickle

from torch.utils.data import TensorDataset
from tqdm import tqdm

logger = logging.getLogger(__name__)

class InputExample(object):
    """A single training/test example for simple sequence classification."""

    def __init__(self, guid, text, label=None):
        self.guid = guid
        self.text = text 
        self.label = label

class InputFeatures(object):
    """A single set of features of data."""

    def __init__(self, input_ids, input_mask, segment_ids, label_id, ori_tokens):

        self.input_ids = input_ids
        self.input_mask = input_mask
        self.segment_ids = segment_ids
        self.label_id = label_id
        self.ori_tokens = ori_tokens


class NerProcessor(object):
    def read_data(self, input_file):
        """Reads a BIO data."""
        with open(input_file, "r", encoding="utf-8") as f:
            lines = []
            words = []
            labels = []
            for line in f.readlines():   
                contends = line.strip()
                tokens = line.strip().split()
                if len(tokens) == 2:
                    words.append(tokens[0])
                    labels.append(tokens[1])
                else:
                    if len(contends) == 0 and len(words) > 0:
                        label = []
                        word = []
                        for l, w in zip(labels, words):
                            if len(l) > 0 and len(w) > 0:
                                label.append(l)
                                word.append(w)
                        lines.append([' '.join(label), ' '.join(word)])
                        words = []
                        labels = []
            
            return lines
            '''
            lines:   [['O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O B-V E-V O O O B-N E-N O O O O', '从 这 些 反 应 中 ， 我 们 清 楚 地 看 出 农 民 们 为 了 急 于 提 高 自 己 的 文 化 水 平 ， 对 “ 速 成 识 字 法 ” 寄 予 了 很 大 希 望 ， 深 望 它']]

            '''
    
    def get_labels(self, args):
        labels = set()
        if os.path.exists(os.path.join(args.output_dir, "label_list.pkl")):
            logger.info(f"loading labels info from {args.output_dir}")
            with open(os.path.join(args.output_dir, "label_list.pkl"), "rb") as f:
                labels = pickle.load(f)
        else:
            # get labels from train data
            logger.info(f"loading labels info from train file and dump in {args.output_dir}")
            with open(args.train_file) as f:
                for line in f.readlines():
                    tokens = line.strip().split()

                    if len(tokens) == 2:
                        labels.add(tokens[1])

            if len(labels) > 0:
                with open(os.path.join(args.output_dir, "label_list.pkl"), "wb") as f:
                    pickle.dump(labels, f)
            else:
                logger.info("loading error and return the default labels B,I,O")
                labels = {"O", "B", "I"}
        
        return labels 

    def get_examples(self, input_file):
        examples = []
        
        lines = self.read_data(input_file)

        for i, line in enumerate(lines):
            guid = str(i)
            text = line[1]#"今天天气不错"
            label = line[0]#"O O B I O O"

            examples.append(InputExample(guid=guid, text=text, label=label))
        #examples=[class1.text class1.label, class2, class3, ... , classn]
        return examples


def convert_examples_to_features(args, examples, label_list, max_seq_length, tokenizer):

    label_map = {label : i for i, label in enumerate(label_list)}#label2id
    #{'E-V': 0, 'B-N': 1, 'O': 2, 'E-N': 3, 'I-N': 4, 'B-V': 5}
    features = []

    for (ex_index, example) in tqdm(enumerate(examples), desc="convert examples"):#这是在实现给word打标签,一一对应label,子单词
        # if ex_index % 10000 == 0:
        #     logger.info("Writing example %d of %d" % (ex_index, len(examples)))
        
        textlist = example.text.split(" ")#['从', '这', '些', '反', '应', '中']
        labellist = example.label.split(" ")
        assert len(textlist) == len(labellist)
        tokens = []
        labels = []
        ori_tokens = []

        for i, word in enumerate(textlist):
            # 防止wordPiece情况出现，不过貌似不会
            
            token = tokenizer.tokenize(word)
            tokens.extend(token)
            label_1 = labellist[i]
            ori_tokens.append(word)

            # 单个字符不会出现wordPiece
            for m in range(len(token)):
                if m == 0:
                    labels.append(label_1)
                else:
                    if label_1 == "O":
                        labels.append("O")
                    else:#为 B 或者 I
                        labels.append("I")#如果有子单词，id加载为I
            
        if len(tokens) >= max_seq_length - 1:
            tokens = tokens[0:(max_seq_length - 2)]  # -2 的原因是因为序列需要加一个句首和句尾标志
            labels = labels[0:(max_seq_length - 2)]
            ori_tokens = ori_tokens[0:(max_seq_length - 2)]

        ori_tokens = ["[CLS]"] + ori_tokens + ["[SEP]"]
        
        ntokens = []
        segment_ids = []
        label_ids = []
        ntokens.append("[CLS]")
        segment_ids.append(0)
        label_ids.append(label_map["O"])#开始的CLS对应label指定为：O

        for i, token in enumerate(tokens):
            ntokens.append(token)
            segment_ids.append(0)
            label_ids.append(label_map[labels[i]])

        ntokens.append("[SEP]")
        segment_ids.append(0)
        label_ids.append(label_map["O"])#最后加的是CLS='O'
        input_ids = tokenizer.convert_tokens_to_ids(ntokens)   
        
        input_mask = [1] * len(input_ids)

        assert len(ori_tokens) == len(ntokens), f"{len(ori_tokens)}, {len(ntokens)}, {ori_tokens}"

        while len(input_ids) < max_seq_length:
            input_ids.append(0)
            input_mask.append(0)
            segment_ids.append(0)
            # we don't concerned about it!
            label_ids.append(0)#padding直接补零
            ntokens.append("**NULL**")

        assert len(input_ids) == max_seq_length
        assert len(input_mask) == max_seq_length
        assert len(segment_ids) == max_seq_length
        assert len(label_ids) == max_seq_length

        if ex_index < 5:
            logger.info("*** Example ***")
            logger.info("guid: %s" % (example.guid))
            logger.info("tokens: %s" % " ".join(
                [str(x) for x in ntokens]))
            logger.info("input_ids: %s" % " ".join([str(x) for x in input_ids]))
            logger.info("input_mask: %s" % " ".join([str(x) for x in input_mask]))
            logger.info("segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
            logger.info("label_ids: %s" % " ".join([str(x) for x in label_ids]))

        # if not os.path.exists(os.path.join(output_dir, 'label2id.pkl')):
        #     with open(os.path.join(output_dir, 'label2id.pkl'), 'wb') as w:
        #         pickle.dump(label_map, w)

        features.append(
                InputFeatures(input_ids=input_ids,
                              input_mask=input_mask,
                              segment_ids=segment_ids,
                              label_id=label_ids,
                              ori_tokens=ori_tokens))#原始句子token,以SEP结尾,没有加padding

    return features


def get_Dataset(args, processor, tokenizer, mode="train"):
    if mode == "train":
        filepath = args.train_file
    elif mode == "eval":
        filepath = args.eval_file
    elif mode == "test":
        filepath = args.test_file
    else:
        raise ValueError("mode must be one of train, eval, or test")

    examples = processor.get_examples(filepath)#[class1[.text,.label],class2...classn]
    label_list = args.label_list

    features = convert_examples_to_features(
        args, examples, label_list, args.max_seq_length, tokenizer
    )

    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
    all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)
    all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)
    all_label_ids = torch.tensor([f.label_id for f in features], dtype=torch.long)

    data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)

    return examples, features, data



###简化版本实现 convert_examples_to_features：
'''
textlist=['实', '际', '上', '，', '不', '仅', '法', '国', '殖', '民', '主', '义', '者', '，', '而', '且', '美', '、', '英', '殖', '民', '主', '义', '者', '都', '对', '这', '个', '“', '根', '本', '法', '”', '草', '案', '寄', '予', '很', '大', '希', '望', '。']
labellist=['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'B-V', 'E-V', 'O', 'O', 'B-N', 'E-N', 'O']
label_map={'B-V': 0, 'E-N': 1, 'E-V': 2, 'B-N': 3, 'O': 4}
print("label_map: ",label_map)

from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
tokenized_input = tokenizer(textlist, is_split_into_words=True)

print("tokenized_input: ",tokenized_input)
ntokens = tokenizer.convert_ids_to_tokens(tokenized_input["input_ids"])
print("ntokens: ",ntokens)
label_ids=[]
label_ids.append(label_map["O"])#开始加的是CLS='O'
for item in labellist:
    label_ids.append(label_map[item])
label_ids.append(label_map["O"])#最后加的是CLS='O'

while len(tokenized_input["input_ids"]) < 50:
    tokenized_input["input_ids"].append(0)
    tokenized_input["attention_mask"].append(0)
    tokenized_input["token_type_ids"].append(0)
    # we don't concerned about it!
    label_ids.append(0)#padding直接补零
    ntokens.append("**NULL**")
print("input_ids: ",tokenized_input["input_ids"])
print("input_mask: ",tokenized_input["attention_mask"])
print("segment_ids: ",tokenized_input["token_type_ids"])
print("label_ids: ",label_ids)
print("ntokens: ",ntokens)
'''

'''
textlist=['实', '际', '上', '，', '不', '仅', '法', '国', '殖', '民', '主', '义', '者', '，', '而', '且', '美', '、', '英', '殖', '民', '主', '义', '者', '都', '对', '这', '个', '“', '根', '本', '法', '”', '草', '案', '寄', '予', '很', '大', '希', '望', '。']
labellist=['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'B-V', 'E-V', 'O', 'O', 'B-N', 'E-N', 'O']
label_map={'B-V': 0, 'E-N': 1, 'E-V': 2, 'B-N': 3, 'O': 4}
print("label_map: ",label_map)

import transformers
assert isinstance(tokenizer, transformers.PreTrainedTokenizerFast)

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
tokenized_input = tokenizer(textlist, is_split_into_words=True)

print("tokenized_input: ",tokenized_input)
ntokens = tokenizer.convert_ids_to_tokens(tokenized_input["input_ids"])
print("ntokens: ",ntokens)
word_ids=tokenized_input.word_ids()
print("word_ids: ",word_ids)
ll=[]
for item in word_ids:#词的id 索引
    if item is None:
        ll.append(label_map["O"])
    else:
        ll.append(label_map[labellist[item]])
print(ll)
while len(tokenized_input["input_ids"]) < 50:
    tokenized_input["input_ids"].append(0)
    tokenized_input["attention_mask"].append(0)
    tokenized_input["token_type_ids"].append(0)
    # we don't concerned about it!
    ll.append(0)#padding直接补零
    ntokens.append("**NULL**")
print("input_ids: ",tokenized_input["input_ids"])
print("input_mask: ",tokenized_input["attention_mask"])
print("segment_ids: ",tokenized_input["token_type_ids"])
print("label_ids: ",ll)
print("ntokens: ",ntokens)
'''